"""
sequence of events
1. Production planning: consumption and capital firms compute their desired output level.

"""
from agent import Firms, HouseHolds, Government, Banks, CentralBank
from typing import List

import pandas as pd
from tqdm import tqdm

from utils import p, INI_YEAR, FirmInfo, BankInfo
from utils import ensure_folder_exists

from agent import Firms, HouseHolds, Government, CentralBank, Banks
from agent import CapitalFirmProduceAndRecruit
from agent import ConsumptionFirmProduceAndRecruit

from contract import CashAdvanceContract
from contract import ReserveContract
from contract import TreasuryDepositContract
from contract import CapitalGoods
from contract import ConsumptionGoods

from market import CreditMarkets
from episode import Events

from monitor import TransactionFlowMonitor
from monitor import BalanceSheetMonitor

pd.set_option('display.max_rows', None)  # 设置显示所有行
pd.set_option('display.max_columns', 11)  # 设置显示所有列
pd.set_option('display.width', None)


def ini_household_agents() -> List:
    """
    初始化居民群体，分配其id（household_id），工资（wages）, 存款(deposits), 近期就业状态(employ_history)
    household_id:分配原则：按照序号累加,比如h_0, h_1
    wages:分配原则：初始化情况下工资相等,
    deposits:存款平均分，总存款 / 居民个体数
    :return:
    """
    n_o_h = int(p.w.number_of_households)
    print(f">>> 初始化居民群体,编号ID from household_0 to household_{n_o_h - 1}")
    households = []
    for h_i in range(n_o_h):
        household_id = f"household_{h_i}"  # 身份编码
        household = HouseHolds(ini_id=household_id)
        households.append(household)
        h_i += 1
    # 后续还需要确定，居民的存款银行，居民的消费企业，居民持有公司，银行的股份
    return households


def ini_firms_agents(firm_type) -> List:
    """
    初始公司基本信息：如id（firm_id）,公司类型（firm_type）,
    存货数量（ini_inventory),单位成本（unit_cost),
    存款（deposits）,贷款（loans）
    firm_id: firm_c_0, firm_c_1 ,,,firm_k_0
    firm_type:两种公司，一种是消费品公司Consumer Goods（c） 另一种是 Capital Goods（k）
    deposits: 公司群体持有存款总数 / 公司个数 （c和k）
    loans:    公司群体贷款总数 / 公司个数 （c和k）
    amortization: c公司摊销奏疏 / c公司个数
    unit_cost : 提前定义
    :param firm_type:
    :return:
    """
    if firm_type == 'Consumer Goods':
        number_of_c_k_firms = p.n.number_of_c_firms
    else:
        number_of_c_k_firms = p.n.number_of_k_firms
    n_c_k = int(number_of_c_k_firms)

    if firm_type == 'Consumer Goods':
        number_of_firms = p.n.number_of_c_firms
        all_firm_loans = p.c.loans_of_consumption_firms
        loans_of_ck = p.c.loans_of_consumption_firms / p.n.number_of_c_firms
        unit_cost_of_ck = p.c.unit_cost_of_consumption_firm
    else:
        number_of_firms = p.n.number_of_k_firms
        all_firm_loans = p.k.loans_of_capital_firms
        loans_of_ck = p.k.loans_of_capital_firms / p.n.number_of_k_firms
        unit_cost_of_ck = p.k.unit_cost_of_capital_firm
    print(f">>> 初始化公司Agent，公司类型为: {firm_type} 为每个群体分配id，数量为:{n_c_k} 编号ID from 0 to {n_c_k - 1} "
          f"每个公司贷款为 {loans_of_ck:.1f}({all_firm_loans:.1f}/{number_of_firms})")

    firms = []
    for ck_i in range(number_of_c_k_firms):
        # 消费品公司参数初始化
        expected_sales_last_period = 0
        expected_wages_last_period = 0
        wages_last_period = 0
        sales_last_period = 0
        
        if firm_type == 'Consumer Goods':
            firm_id = f"firm_c_{ck_i}"
            inv = p.c.consumption_firm_inventory / p.n.number_of_c_firms
            inventory = ConsumptionGoods(
                year=INI_YEAR, 
                quantity=inv, 
                unit_cost=unit_cost_of_ck,
                price=unit_cost_of_ck,
                )
            price_and_recruit = ConsumptionFirmProduceAndRecruit(
                year=INI_YEAR,
                sales_last_period=p.y.output_of_consumption / p.n.number_of_c_firms,
                expected_sales_last_period=expected_sales_last_period,
                inventory_last_period_end=inventory.quantity,
                capital_holds=[],
                expected_wages_last_period=expected_wages_last_period,
                wages_last_period=wages_last_period,
                mark_up=p.m.mark_up_of_consumption,
                price_last_period=p.p.price_of_consumption,
                )
        # 资本品公司参数初始化
        else:
            firm_id = f"firm_k_{ck_i}"
            inv = p.k.capital_firm_inventory / p.n.number_of_k_firms
            capital = 0
            inventory = CapitalGoods(
                year=INI_YEAR, 
                quantity=inv, 
                price=unit_cost_of_ck, 
                share=1)
            price_and_recruit = CapitalFirmProduceAndRecruit(
                year=INI_YEAR,
                sales_last_period=sales_last_period,
                expected_sales_last_period=expected_sales_last_period,
                inventory_last_period_end=inventory.quantity,
                expected_wages_last_period=expected_wages_last_period,
                wages_last_period=wages_last_period,
                mark_up=p.m.mark_up_of_capital,
                price_last_period=p.p.price_of_capital,
            )

        price_and_recruit.price = price_and_recruit.this_year_price
        firm = Firms(
            ini_id=firm_id,
            firm_type=firm_type,
            ini_inventory=inv,
            ini_loans=loans_of_ck,
            price_and_recruit=price_and_recruit,
            inventory=inventory
            )
        firms.append(firm)
    return firms


def ini_bank_agents() -> List:
    """
    初始化银行agent，包括id，准备金，贷款，银行持有债券，存款，银行向中央银行的预借款
    :return:
    """
    n_o_b = int(p.n.number_of_banks)  # 银行数量
    print(f">>> 初始化银行Agent，为每个银行分配id，数量为{n_o_b} 编号ID from 0 to {n_o_b - 1}")

    total_loans = p.c.loans_of_consumption_firms + p.k.loans_of_capital_firms
    loans = total_loans / n_o_b
    print(f"  > 【银行初始化】总贷款（{total_loans:.1f}） = 消费品公司贷款({p.c.loans_of_consumption_firms:.1f}) + "
          f"投资品公司贷款({p.k.loans_of_capital_firms:.1f})")
    print(f"  > 【银行初始化】单个银行贷款（{loans:.1f}）= 总贷款({total_loans:.1f}) / 银行数量({n_o_b:.1f})")
    print(f"  > 【银行初始化】单个银行存在中央银行的保证金:({p.b.bank_reserves:.1f})")
    print(f"  > 【银行初始化】单个银行持有国债:({p.b.bank_hold_bonds:.1f})")

    banks = []
    for i_b in range(n_o_b):
        bank_agent = Banks(
            ini_id=f"bank_{i_b}",
            ini_bonds=p.b.bank_hold_bonds / n_o_b,  # 银行持有的国债平均分配到每个银行
            cash_advance_contract=CashAdvanceContract(
                {'amount': 0, 'bank_id': f"bank_{i_b}", 'interest_rate': p.ir.interest_rate_of_advances}),
            reserve_contract=ReserveContract(
                {'amount': p.b.bank_reserves / n_o_b, 'bank_id': f"bank_{i_b}"})
        )
        banks.append(bank_agent)
    return banks


def export_log(e: Events, i: int):
    dic = {
        'h': [e.credit_market.households, '居民日志'],
        'k': [e.credit_market.k_firms, 'K公司日志'],
        'c': [e.credit_market.c_firms, "C公司日志"],
        'b': [e.credit_market.banks, "银行日志"]
    }
    dir_path = f'../log/{i}/'
    ensure_folder_exists(dir_path)
    for k, v in dic.items():
        logs = []
        for obj_i, obj_ in enumerate(v[0]):
            logs.extend(obj_.log)
            if obj_i == 0:
                dir_path_obj = f'../log/{i}/{k}/'
                ensure_folder_exists(dir_path_obj)
                df_obj_log = pd.DataFrame(obj_.log)
                df_obj_log.to_excel(f"{dir_path_obj}{obj_.id_()}.xlsx")
        df = pd.DataFrame(logs)
        print(k, v[-1], df.shape)
        agent_file = f"{dir_path}{v[1]}/"
        ensure_folder_exists(agent_file)
        df.to_excel(f"{agent_file}{v[1]}_{i}总表.xlsx")
        for tf_item in df['transaction_flow_item'].unique():
            df_tf = df[df['transaction_flow_item'] == tf_item]
            df_tf.to_excel(f"{agent_file}{tf_item}_{i}.xlsx")
        # df.to_csv(f"{dir_path}{v[1]}_{i}.csv")
    bs = BalanceSheetMonitor(e) # 计算并展示资产负债表
    df_balance_sheet = bs.show_bs_table(prt=False)
    df_balance_sheet.to_csv(f"{dir_path}资产负债表_{i}.csv")

    tf = TransactionFlowMonitor(eve) # 计算并展示交易流量表
    df_cash_flow = tf.show_tf_table(prt=False)
    df_cash_flow.to_csv(f"{dir_path}资金流量表_{i}.csv")


def save_bs_tf(save_name, record):
    file_path = f"../{save_name}_all_period.xlsx"
    with pd.ExcelWriter(file_path) as writer:
        for k, v in record.items():
            v.to_excel(writer, sheet_name=f"{k}")


if __name__ == "__main__":
    print(">>> 开始")
    # 初始化居民部门
    all_households = ini_household_agents()
    # 初始化企业部门
    all_k_firms = ini_firms_agents(firm_type='Capital Goods')
    all_c_firms = ini_firms_agents(firm_type='Consumer Goods')
    # 初始化银行
    all_banks = ini_bank_agents()
    gov_deposit = TreasuryDepositContract()  # 初始化财政（国库）存款合约
    gov = Government(gov_deposit)  # 初始化政府
    # 初始化中央银行
    cb = CentralBank(
        all_banks=all_banks,
        treasury_deposit_contract=gov_deposit
        )
    # 初始化市场
    credit_market = CreditMarkets(
        year=INI_YEAR,
        c_firms=all_c_firms,
        k_firms=all_k_firms,
        banks=all_banks,
        households=all_households,
        government=gov,
        central_bank=cb,
        )
    # 遍历程序
    eve = Events(
        government=gov,
        central_bank=cb,
        credit_markets=credit_market,
        year=INI_YEAR
        ) 
    
    bs_before_period_record = {}
    bs_after_period_record = {}
    transaction_flow_record = {}

    for i in tqdm(range(400)): # 经济循环迭代次数
        # 期初看看资产负债情况
        bs = BalanceSheetMonitor(eve) # 计算并展示资产负债表(期初)
        bs_before_period_record[eve.year] = bs.show_bs_table()

        eve.run_a_episode() # 运行一个经济周期

        # 看看资金流量情况
        tf = TransactionFlowMonitor(eve) # 计算并展示交易流量表
        tf.check_all_agents_consistent()
        transaction_flow_record[eve.year] = tf.show_tf_table()
        # 期末查看资资产负债表
        # bs = BalanceSheetMonitor(eve) # 计算并展示资产负债表(期末)
        # bs_after_period_record[eve.year] = bs.show_bs_table()

        if abs(tf.k_firm_inv_revalue) > 700:
            print(666)
            k_logs = tf.logs['k']

        eve.year += 1
        eve.credit_market.year += 1
        eve.credit_market.central_bank.year += 1
        eve.credit_market.government.year += 1
        for firms in [eve.credit_market.c_firms, eve.credit_market.k_firms]:
            for f in firms:
                f.year += 1
        for household in eve.credit_market.households:
            household.year += 1
        for b in eve.credit_market.banks:
            b.year += 1
        
        if i % 40 == 0:
            save_bs_tf('资产负债表', bs_before_period_record)
            save_bs_tf('资金流量表', transaction_flow_record)

        print(">>"*36, f"【{i}】", "<<"*36)

print(f"==={INI_YEAR}年经济循环结束==",f"==={INI_YEAR}年经济循环结束==")